from types import SimpleNamespace

CONFIG_ZIPFORMER = {
    "train_switch": "u2",
    "world_size": 7,
    "master_port": 12368,
    "tensorboard": True,
    "num_epochs": 10,
    "start_epoch": 4,
    "start_batch": 0,
    # "exp_dir": "zipformer_streaming_u2/exp-u2-stage04",
    "lang_dir": "data/lang_char_bpe",
    "bpe_model": "data/lang/bpe.model",
    "base_lr": 0.025,
    "lr_batches": 7500,
    "lr_epochs": 3.5,
    "ref_duration": 600,
    "context_size": 2,
    "prune_range": 5,
    "checkpoint_path": "no_checkpoint",
    "lm_scale": 0.25,
    "am_scale": 0.0,
    "simple_loss_scale": 0.5,
    "ctc_loss_scale": 0.1,
    "attention_decoder_loss_scale": 0.9,
    "seed": 42,
    "print_diagnostics": False,
    "inf_check": False,
    "save_every_n": 10000,
    "keep_last_k": 30,
    "average_period": 200,
    "use_fp16": True,
    "use_tal_csasr": False,
    "use_librispeech": False,
    "use_aishell2": False,
    "num_encoder_layers": '2, 2, 4, 5, 4, 2',
    "downsampling_factor": '1, 2, 4, 8, 4, 2',
    "feedforward_dim": '512, 768, 1536, 2048, 1536, 768',
    "num_heads": '4, 4, 4, 8, 4, 4',
    "encoder_dim": '192, 256, 512, 768, 512, 256',
    "query_head_dim": 32,
    "value_head_dim": 12,
    "pos_head_dim": 4,
    "pos_dim": 48,
    "encoder_unmasked_dim": '192, 192, 256, 320, 256, 192',
    "cnn_module_kernel": '31, 31, 15, 15, 15, 31',
    "decoder_dim": 512,
    "joiner_dim": 512,
    "attention_decoder_dim": 512,
    "attention_decoder_num_layers": 6,
    "attention_decoder_attention_dim": 512,
    "attention_decoder_num_heads": 8,
    "attention_decoder_feedforward_dim": 2048,
    "causal": True,
    # "chunk_size": '16, 32, 64, -1',
    "chunk_size": '32', # 导出的时候由导出文件提供该数值
    # "left_context_frames": '64, 128, 256, -1',
    "left_context_frames": '1024',
    "use_transducer": True,
    "use_ctc": True,
    "use_attention_decoder": True,
    "valid_fbank_dir": "/home/local_data/yhdai/workspace/icefall/egs/multi_zh_en/ASR/data/dev_test_cuts/dev_family",
    "valid_prefix": "valid_01",
    "max_duration": 600,
    "bucketing_sampler": True,
    "num_buckets": 150,
    "concatenate_cuts": False,
    "duration_factor": 1.0,
    "gap": 1.0,
    "on_the_fly_feats": False,
    "shuffle": True,
    "drop_last": True,
    "return_cuts": True,
    "num_workers": 5,
    "enable_spec_aug": True,
    "spec_aug_time_warp_factor": 80,
    "enable_musan": True,
    "input_strategy": "PrecomputedFeatures",
    "blank_id": 0,
    "sos_id": 1,
    "eos_id": 1,
    "vocab_size": 6254,
    "cur_epoch": 10,
    "train_loss": 0.4759614318019831,
    "log_interval": 50,
    "reset_interval": 200,
    "valid_interval": 3000,
    "feature_dim": 80,
    "subsampling_factor": 4,
    "ignore_id": -1,
    "label_smoothing": 0.1,
    "warm_step": 2000
}

CONFIG_ZIPFORMER_PARAM = SimpleNamespace(**CONFIG_ZIPFORMER)